"""
Pinecone MNIST训练数据上传与准确率测试工具
依赖安装: pip install pinecone scikit-learn numpy tqdm

使用说明:
- 需先注册Pinecone并获取API密钥
- 免费用户请使用us-east-1区域的serverless配置
- 数据集划分: 训练集约1437条，测试集约360条
"""
from __future__ import annotations
import time
import logging
from typing import List
import numpy as np
from tqdm import tqdm
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from pinecone import Pinecone, ServerlessSpec

# 配置参数
INDEX_TITLE = "mnist-index"
VECTOR_LENGTH = 64
DISTANCE_METRIC = "euclidean"
BATCH_UPLOAD_SIZE = 100
API_KEY = "pcsk_4EcLgW_MpufGtAtj64LC3kASF1U834xnJP1M12k6ihSTyfkBteSHcGQpdTmjLVZzMkPRdb"


def configure_logging() -> None:
    """配置日志输出格式"""
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s,%(msecs)03d %(levelname)s %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )


def manage_index(pc: Pinecone, idx_name: str) -> None:
    """管理索引：存在则删除并重新创建"""
    try:
        existing = set(pc.list_indexes().names())
    except Exception:
        try:
            existing = {i["name"] for i in pc.list_indexes()}
        except Exception:
            existing = set()

    if idx_name in existing:
        logging.info("索引 %s 已存在，正在删除...", idx_name)
        pc.delete_index(idx_name)
        time.sleep(2)

    logging.info("创建索引 %s (维度=%d, 度量=%s, 区域=us-east-1)...", 
                 idx_name, VECTOR_LENGTH, DISTANCE_METRIC)
    pc.create_index(
        name=idx_name,
        dimension=VECTOR_LENGTH,
        metric=DISTANCE_METRIC,
        spec=ServerlessSpec(cloud="aws", region="us-east-1")
    )
    time.sleep(5)  # 等待索引就绪
    logging.info("索引创建完成")


def convert_to_vectors(X: np.ndarray, y: np.ndarray) -> List[dict]:
    """将数据集转换为Pinecone向量格式"""
    vector_items = []
    for i in range(len(X)):
        vector_items.append({
            "id": str(i),
            "values": X[i].astype(float).tolist(),
            "metadata": {"Label": int(y[i])}
        })
    return vector_items


def batch_upload(index, vectors: List[dict], batch_size: int = BATCH_UPLOAD_SIZE) -> None:
    """带进度条的批量上传"""
    total = len(vectors)
    batches = (total + batch_size - 1) // batch_size
    for batch_num in tqdm(range(batches), desc="上传至Pinecone", unit="批"):
        start = batch_num * batch_size
        end = min(start + batch_size, total)
        index.upsert(vectors=vectors[start:end])


def evaluate_accuracy(index, X_test: np.ndarray, y_test: np.ndarray, k_neighbors: int = 11) -> float:
    """评估KNN分类准确率"""
    predictions = []
    for i in tqdm(range(len(X_test)), desc=f"测试k={k_neighbors}准确率", unit="样本"):
        vector = X_test[i].astype(float).tolist()
        result = index.query(vector=vector, top_k=k_neighbors, include_metadata=True)
        matches = result.get("matches", []) if isinstance(result, dict) else getattr(result, "matches", [])
        
        # 收集邻居标签
        neighbor_labels = []
        for match in matches:
            meta = match.get("metadata") if isinstance(match, dict) else getattr(match, "metadata", {})
            if "Label" in meta:
                neighbor_labels.append(int(meta["Label"]))
        
        # 投票决定预测结果
        if neighbor_labels:
            pred = int(np.argmax(np.bincount(neighbor_labels)))
        else:
            pred = 0  # 默认值
        predictions.append(pred)
    
    return float(accuracy_score(y_test, predictions))


def execute_training() -> None:
    """主执行函数"""
    configure_logging()

    # 加载并划分数据集
    digits = load_digits(n_class=10)
    X = digits.data.astype(np.float32)
    y = digits.target.astype(int)

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y
    )
    logging.info("训练集大小: %d, 测试集大小: %d", len(X_train), len(X_test))

    # 初始化Pinecone
    try:
        pinecone_client = Pinecone(api_key=API_KEY)
    except Exception as e:
        err_msg = str(e)
        if "Invalid API Key" in err_msg or "401" in err_msg:
            logging.error("API密钥无效，请检查")
        else:
            logging.error("初始化失败: %s", err_msg)
        return

    # 管理索引
    manage_index(pinecone_client, INDEX_TITLE)
    index = pinecone_client.Index(INDEX_TITLE)

    # 上传训练数据
    train_vectors = convert_to_vectors(X_train, y_train)
    batch_upload(index, train_vectors)
    logging.info("成功上传 %d 条训练数据", len(train_vectors))

    # 评估准确率
    accuracy = evaluate_accuracy(index, X_test, y_test, k_neighbors=11)
    logging.info("k=11时，Pinecone KNN准确率: %.3f", accuracy)


if __name__ == "__main__":
    execute_training()